import os
import re
import sys
import argparse
import subprocess
from pathlib import Path
from typing import Dict, Optional, Tuple

try:
    from openpyxl import load_workbook  # Optional, used for Excel fallback
except Exception:
    load_workbook = None

# 分类标签（保持与 config/categories_example.txt 一致，全角符号）
DURATION_LABELS = ["≤15秒", "≤1分钟", "≤5分钟", "＞5分钟"]


def classify_duration(seconds: int) -> str:
    if seconds <= 15:
        return DURATION_LABELS[0]
    if seconds <= 60:
        return DURATION_LABELS[1]
    if seconds <= 300:
        return DURATION_LABELS[2]
    return DURATION_LABELS[3]


def parse_txt_filename(fname: str) -> Tuple[Optional[str], Optional[str], Optional[str]]:
    """解析 TXT 文件名，提取时间戳、作者、标题（尽力而为）。
    期望格式：YYYY-MM-DD HH.MM.SS-视频-作者-标题.txt
    返回：(timestamp, author, title)
    """
    base = fname[:-4] if fname.lower().endswith('.txt') else fname
    # 切分时保留中文与空格
    parts = base.split('-')
    if len(parts) < 3:
        return None, None, None
    timestamp = parts[0].strip()
    # 期望第二段是“视频”
    author = None
    title = None
    if len(parts) >= 3 and parts[1] == '视频':
        author = parts[2].strip()
        title = '-'.join(parts[3:]).strip() if len(parts) > 3 else None
    else:
        # 退化处理：第二段不是“视频”，则尝试把第二段当作者
        author = parts[1].strip()
        title = '-'.join(parts[2:]).strip() if len(parts) > 2 else None
    return timestamp, author, title


def normalize_title(title: str) -> str:
    if not title:
        return ''
    # 去空白与标点（含中英文常见标点，保留汉字与数字字母）
    t = re.sub(r"\s+", "", title)
    t = re.sub(r"[，。！？；：、（）《》“”'\"\-_#@\[\]()…]+", "", t)
    return t.lower()


def discover_video_path(video_root: Path, base_name: str, lazy_index: Dict[str, Path]) -> Optional[Path]:
    """根据 TXT 的 base_name，寻找对应视频路径。
    优先尝试 video_root 下的同名 .mp4；未命中时懒构建索引（首次全量扫描）。
    """
    direct = video_root / f"{base_name}.mp4"
    if direct.exists():
        return direct
    # 懒索引：只在需要时构建一次
    if not lazy_index:
        for p in video_root.rglob('*'):
            if p.is_file() and p.suffix.lower() in {'.mp4', '.mov', '.mkv'}:
                lazy_index[p.stem] = p
    return lazy_index.get(base_name)


def ffprobe_duration(video_path: Path) -> Optional[int]:
    """用 ffprobe 获取时长（秒，取整）。"""
    try:
        # Windows PowerShell 环境下同样可执行
        cmd = [
            'ffprobe', '-v', 'error',
            '-show_entries', 'format=duration',
            '-of', 'default=noprint_wrappers=1:nokey=1',
            str(video_path)
        ]
        out = subprocess.check_output(cmd, stderr=subprocess.STDOUT, timeout=15)
        txt = out.decode('utf-8', errors='ignore').strip()
        if not txt:
            return None
        # 某些容器可能返回浮点秒数
        val = float(txt)
        sec = int(round(val))
        if sec <= 0:
            return None
        return sec
    except Exception:
        return None


def parse_duration_text(v: str) -> Optional[int]:
    """解析 Excel 中的时长文本，支持：
    - 纯数字秒（如 160）
    - 浮点秒（如 160.2）
    - mm:ss（如 02:40）
    - 含单位（如 "160秒"、"2分40秒"、"2m40s"）
    """
    if v is None:
        return None
    if isinstance(v, (int, float)):
        sec = int(round(v))
        return sec if sec > 0 else None
    s = str(v).strip()
    if not s:
        return None
    # mm:ss
    m = re.match(r"^(\d{1,3}):(\d{2})$", s)
    if m:
        mm = int(m.group(1))
        ss = int(m.group(2))
        return mm * 60 + ss
    # 纯数字
    if re.fullmatch(r"\d+(\.\d+)?", s):
        return int(round(float(s)))
    # 中文单位与英文单位的宽松解析
    s2 = s.replace('分', 'm').replace('分钟', 'm').replace('秒', 's')
    s2 = s2.replace('：', ':')  # 全角冒号
    # 可能是 like "2m40s"
    m2 = re.match(r"^(\d{1,3})m(\d{1,2})s$", s2)
    if m2:
        return int(m2.group(1)) * 60 + int(m2.group(2))
    # 可能是 like "160s"
    m3 = re.match(r"^(\d{1,5})s$", s2)
    if m3:
        return int(m3.group(1))
    return None


def load_excel_map(excel_path: Path,
                    duration_col_idx: int = 6,
                    timestamp_col_idx: Optional[int] = 1,
                    author_col_idx: Optional[int] = 3,
                    title_col_idx: Optional[int] = 4) -> Dict[str, int]:
    """读取 Excel，构建兜底字典。
    列索引从 1 开始。默认 F 列为时长（6）。其余列若不可用则降级。
    key 采用两种形式，优先匹配 TS（含时间戳）：
      - TS:{timestamp}|{author}|{normalized_title[:30]}
      - AT:{author}|{normalized_title[:30]}（当无法可靠获取时间戳时作为次级兜底）
    """
    mapping: Dict[str, int] = {}
    if not load_workbook:
        return mapping
    if not excel_path.exists():
        return mapping
    try:
        wb = load_workbook(filename=str(excel_path), read_only=True, data_only=True)
        ws = wb.active
        for i, row in enumerate(ws.iter_rows(values_only=True), start=1):
            try:
                duration_val = row[duration_col_idx - 1] if len(row) >= duration_col_idx else None
                sec = parse_duration_text(duration_val)
                if not sec:
                    continue
                ts = None
                au = None
                ti = None
                if timestamp_col_idx and len(row) >= timestamp_col_idx:
                    ts_raw = row[timestamp_col_idx - 1]
                    ts = str(ts_raw).strip() if ts_raw is not None else None
                if author_col_idx and len(row) >= author_col_idx:
                    au_raw = row[author_col_idx - 1]
                    au = str(au_raw).strip() if au_raw is not None else None
                if title_col_idx and len(row) >= title_col_idx:
                    ti_raw = row[title_col_idx - 1]
                    ti = str(ti_raw).strip() if ti_raw is not None else None
                norm_title = normalize_title(ti)[:30] if ti else ''
                # 构建 TS 与 AT 两类键
                if ts and au:
                    key_ts = f"TS:{ts}|{au}|{norm_title}"
                    if key_ts not in mapping:
                        mapping[key_ts] = sec
                if au and ti:
                    key_at = f"AT:{au}|{norm_title}"
                    if key_at not in mapping:
                        mapping[key_at] = sec
            except Exception:
                continue
        return mapping
    except Exception:
        return mapping


def find_duration_from_excel(excel_map: Dict[str, int], timestamp: Optional[str], author: Optional[str], title: Optional[str]) -> Optional[int]:
    if not excel_map or not author:
        return None
    norm = normalize_title(title)[:30] if title else ''
    # 先尝试带时间戳的匹配
    if timestamp:
        val = excel_map.get(f"TS:{timestamp}|{author}|{norm}")
        if val:
            return val
    # 再尝试作者+标题的兜底匹配
    if title:
        return excel_map.get(f"AT:{author}|{norm}")
    return None


def update_txt_duration_line(txt_path: Path, label: str) -> bool:
    try:
        content = txt_path.read_text(encoding='utf-8')
        lines = content.splitlines(keepends=False)
        changed = False
        for idx, line in enumerate(lines):
            if line.startswith('视频时长:'):
                if line.strip() == f"视频时长: {label}":
                    return False  # 无改动
                lines[idx] = f"视频时长: {label}"
                changed = True
                break
        if not changed:
            # 若没有发现该行，则将其插入为第一行（谨慎处理）
            lines.insert(0, f"视频时长: {label}")
            changed = True
        if changed:
            txt_path.write_text("\n".join(lines) + "\n", encoding='utf-8')
        return changed
    except Exception:
        return False


def main():
    parser = argparse.ArgumentParser(description='仅纠正 results/*.txt 的“视频时长”行，按真实时长分档。')
    parser.add_argument('--results-dir', default=str(Path('results').resolve()), help='TXT结果目录')
    parser.add_argument('--video-root', default='S:\\orignV\\', help='原始视频根目录')
    parser.add_argument('--excel-path', default=str(Path('d:/Desktop/DYVA/V2/研究样本表.xlsx')), help='Excel兜底路径')
    parser.add_argument('--excel-duration-col', type=int, default=6, help='Excel时长列（从1开始，默认F=6）')
    parser.add_argument('--excel-timestamp-col', type=int, default=5, help='Excel时间戳列（从1开始，默认5=发布时间）')
    parser.add_argument('--excel-author-col', type=int, default=1, help='Excel作者列（从1开始，默认1=博主名称）')
    parser.add_argument('--excel-title-col', type=int, default=2, help='Excel标题列（从1开始，默认2=视频标题）')

    args = parser.parse_args()
    results_dir = Path(args.results_dir)
    video_root = Path(args.video_root)
    excel_path = Path(args.excel_path)

    # Excel兜底字典（可选）
    excel_map = load_excel_map(
        excel_path,
        duration_col_idx=args.excel_duration_col,
        timestamp_col_idx=args.excel_timestamp_col,
        author_col_idx=args.excel_author_col,
        title_col_idx=args.excel_title_col,
    )

    if not results_dir.exists():
        print(f"[错误] 结果目录不存在: {results_dir}")
        sys.exit(1)

    lazy_index: Dict[str, Path] = {}
    updated, skipped, ffprobe_cnt, excel_cnt = 0, 0, 0, 0

    txt_files = sorted(results_dir.glob('*.txt'))
    print(f"[扫描] 待处理TXT文件: {len(txt_files)}")

    for i, txt_path in enumerate(txt_files, start=1):
        base_name = txt_path.stem
        timestamp, author, title = parse_txt_filename(txt_path.name)
        duration_sec: Optional[int] = None

        # 先 ffprobe
        vpath = discover_video_path(video_root, base_name, lazy_index)
        if vpath:
            duration_sec = ffprobe_duration(vpath)
            if duration_sec:
                ffprobe_cnt += 1
        # 再 Excel 兜底
        if not duration_sec:
            duration_sec = find_duration_from_excel(excel_map, timestamp, author, title)
            if duration_sec:
                excel_cnt += 1

        if not duration_sec:
            skipped += 1
            continue

        label = classify_duration(duration_sec)
        if update_txt_duration_line(txt_path, label):
            updated += 1
        else:
            skipped += 1

        if i % 50 == 0:
            total = updated + skipped
            sr = (updated / total * 100.0) if total else 0.0
            print(f"[进度] {i} / {len(txt_files)} | 更新: {updated} | 跳过: {skipped} | 成功率: {sr:.1f}% | 源: ffprobe={ffprobe_cnt}, excel={excel_cnt}")

    total = updated + skipped
    sr = (updated / total * 100.0) if total else 0.0
    print(f"[完成] 更新: {updated} | 跳过: {skipped} | 成功率: {sr:.1f}% | 源: ffprobe={ffprobe_cnt}, excel={excel_cnt}")


if __name__ == '__main__':
    main()

